import torch
from torch import nn
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import CSVLogger
from pytorch_lightning.callbacks import TQDMProgressBar
from pytorch_lightning.strategies import DDPStrategy

import pandas as pd
import numpy as np
from pathlib import Path
import json
import umap
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
import sys
from torch.utils.data import DataLoader, Sampler, Dataset, DistributedSampler
# --- 调整 sys.path 以便从项目根目录导入 muta 包和 light 包 ---
# 脚本在 src/scripts/ (假设此脚本位于此路径)
# 如果脚本位置不同，请调整 parents[X] 的数字
project_root = Path(__file__).resolve().parents[2] # 向上两级到 主目录
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))
if str(project_root / 'src') not in sys.path:
    sys.path.insert(0, str(project_root / 'src'))

# --- 直接导入 muta 包中的类 ---
from muta.model.mutation_performer import PerformerMutationLM, PerformerMutationClassifier
from muta.data.mutation_data import MutationPretrainDataset, MutationClassificationDataset, read_df_path
from mutation_light import MutationDataModule, get_model, LengthBasedBatchSampler, collate_fn_for_mutation_data
from light.core.random import worker_init_fn_seed # 确保导入这个


# --- 用于提取嵌入的自定义 LightningModule ---
class EmbeddingExtractorModule(pl.LightningModule):
    def __init__(self, task, model_path, data_root):
        super().__init__()
        self.task = task
        self.model_path = Path(model_path)
        self.data_root = Path(data_root)

        print(f"正在从 {self.model_path} 加载模型配置和模型...")

        # 从 model_config.json 读取模型参数
        model_config_path = self.model_path.parent / 'model_config.json'
        if not model_config_path.exists():
            raise FileNotFoundError(f"模型配置文件未找到: {model_config_path}. 无法正确加载模型。")

        with open(model_config_path, "r") as f:
            model_config = json.load(f)

        # 使用 get_model 函数来实例化模型，它会处理所有参数
        self.model = get_model(
            task=model_config.get('task', self.task),
            model_path=self.model_path,
            num_classes=model_config.get('num_classes', 1),
            model_args=model_config
        )

        print("模型加载成功。")
        self.model.eval() # 设置为评估模式，关闭Dropout等

        # 检查加载的模型对象是否包含预期的嵌入层和Performer骨干网络
        # 根据 mutation_performer.py 中的新结构进行检查
        expected_attrs = [
            'pentanucleotide_context_embedding',
            'chromosome_embedding',
            'exon_strand_embedding',
            'sequencing_method_embedding',
            'performer',
            'cls_token', # 确保 cls_token 存在
            'fusion_weights', # 确保 fusion_weights 存在
            'fusion_norm', # 确保 fusion_norm 存在
        ]
        # 根据模型配置动态添加 evo2 和绝对位置编码的属性检查
        if model_config.get('use_absolute_genomic_pos_emb', False):
            expected_attrs.extend(['absolute_genomic_pos_emb', 'absolute_genomic_pos_encoder'])
        if model_config.get('evo2_embedding_dim', 0) > 0:
            expected_attrs.append('evo2_encoder')

        missing_attrs = [attr for attr in expected_attrs if not hasattr(self.model, attr)]
        if missing_attrs:
            # 针对性地警告，而不是直接报错，因为某些属性可能在模型中是可选的
            warnings.warn(f"加载的模型对象缺少以下预期属性: {missing_attrs}. 这可能导致嵌入提取失败。请检查model.pt是否正确。")


        self.all_embeddings = []
        self._test_step_embeddings_buffer = [] # 保持 test_step 的命名，但用于推理

    def forward(self, inputs):
        # 直接调用模型的 forward 方法，它会处理所有特征嵌入、融合和 Performer 传递
        # 假设 PerformerMutationLM 和 PerformerMutationClassifier 的 forward 都返回包含 'sample_embedding' 的字典
        self.model.eval() # 确保在推理时模型处于评估模式
        with torch.no_grad():
            model_outputs = self.model(inputs)
            # 提取 sample_embedding
            if 'sample_embedding' in model_outputs:
                return model_outputs['sample_embedding']
            else:
                raise ValueError(f"模型 {type(self.model)} 的 forward 方法未返回 'sample_embedding'。请检查模型实现。")

    def test_step(self, batch, batch_idx):
        inputs, _ = batch
        embeddings = self(inputs)
        self._test_step_embeddings_buffer.append(embeddings.cpu())

    def on_test_epoch_end(self):
        if self._test_step_embeddings_buffer:
            all_embeddings_tensor = torch.cat(self._test_step_embeddings_buffer, dim=0)
            self.all_embeddings = all_embeddings_tensor.numpy()
            self._test_step_embeddings_buffer.clear()
        else:
            self.all_embeddings = np.array([])
            print("未收集到任何测试步骤输出。")

# --- 封装绘图逻辑的函数 (更新为您的参考版本) ---
def plot_umap(df_data, hue_column, title, output_path, palette=None):
    """
    生成并保存 UMAP 图，具有 1:1 纵横比和单列图例。

    参数:
        df_data (pd.DataFrame): 包含 UMAP 嵌入 ('umap_x', 'umap_y')
                                和用于颜色标记的列 (hue_column) 的 DataFrame。
        hue_column (str): 用于颜色标记点的列名。
        title (str): 图表的标题。
        output_path (str): 保存图表的路径。
        palette (list 或 dict): 用于 hue_column 的颜色调色板。
                                可以是颜色列表，也可以是将 hue 值映射到颜色的字典。
    """
    # 创建一个图形和坐标轴对象。
    # figsize 选择为宽于高，以便为右侧的图例留出空间。
    # 实际绘图区域将通过 ax.set_aspect('equal') 设置为正方形。
    fig, ax = plt.subplots(figsize=(16, 12)) # 示例: 16英寸宽，12英寸高

    if palette is None:
        num_unique_hues = len(df_data[hue_column].unique())
        # 'husl' 通常适用于生成大量感知上不同的颜色。
        # 'tab20' 适用于最多 20 个类别。如果超过，它会循环使用颜色。
        # 对于超过 20-30 个类别，通常首选 'husl'。
        if num_unique_hues <= 20:
            current_palette = 'tab20'
        else:
            current_palette = sns.color_palette("hls", num_unique_hues) # 保持与原代码一致使用hls
    else:
        current_palette = palette

    # 使用 seaborn 的 scatterplot 绘制 UMAP 嵌入
    sns.scatterplot(
        data=df_data,
        x='umap_x',
        y='umap_y',
        hue=hue_column,
        palette=current_palette,
        s=5,           # 标记点的大小，根据数据密度调整
        linewidth=0,   # 标记点无边框，使图表更清晰
        alpha=0.7,     # 透明度，对于密集图表有助于可视化重叠区域
        ax=ax          # 指定绘图的坐标轴
    )

    # 设置图表标题和轴标签
    ax.set_title(title, fontsize=16, pad=20) # pad 增加标题与图表之间的距离
    ax.set_xlabel('UMAP Dimension 1', fontsize=12)
    ax.set_ylabel('UMAP Dimension 2', fontsize=12)

    # 确保绘图区域具有 1:1 的纵横比
    # 'box' 确保限制调整以使数据适应正方形纵横比
    ax.set_aspect('equal', adjustable='box')

    # 自定义图例
    # 获取散点图的图例句柄和标签
    handles, labels = ax.get_legend_handles_labels()

    # 显式创建新图例以控制其属性
    # bbox_to_anchor=(1.02, 1) 将图例框的左上角放置在绘图区域右边缘的 2% 处，并与顶部边缘对齐。
    # loc='upper left' 将图例的左上角与 bbox_to_anchor 对齐。
    # ncol=1 将图例条目排列为单列。
    # fontsize 和 markerscale 可以根据可读性进行调整，特别是当类别很多时。
    ax.legend(
        handles=handles,
        labels=labels,
        bbox_to_anchor=(1.02, 1), # 相对于坐标轴的位置
        loc='upper left',         # 图例框的锚点
        borderaxespad=0.,         # 图例与坐标轴边框之间无填充
        ncol=1,                   # 单列图例
        fontsize='small',         # 字体大小，例如 'x-small', 8, 10
        markerscale=2             # 图例中标记点的大小
    )

    # 移除网格线，使 UMAP 可视化更清晰
    ax.grid(False)

    # 使用 tight_layout 自动调整图表参数以实现紧凑布局。
    # 这通常与 bbox_to_anchor 配合使用，以确保图例包含在内。
    # bbox_inches='tight' 尝试使图表的边界框紧密围绕所有元素，包括放置在坐标轴外部的图例。
    # 这对于确保图表的边界框紧密围绕所有元素，包括放置在坐标轴外部的图例。
    plt.tight_layout()

    # 保存图表，高 DPI。
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.show() # 保持原代码的 plt.show()
    print(f"图已成功保存到: {output_path}")
    plt.close(fig) # 关闭图表以释放内存


# --- 主脚本执行函数 ---
def generate_embeddings_and_umap_plot():
    # 如果要提取预训练模型的嵌入，这里指向预训练模型的 model.pt
    # 如果要提取微调模型的嵌入，这里指向微调模型的 model.pt
    MODEL_PATH = "/gpfs/flash/home/yzq/project/model/mutabert/results/training_output/pretrain_run_gpu_with_evo2_revised_performer_0909/model.pt"
    # MODEL_PATH = "/gpfs/flash/home/yzq/project/model/Mutabert/results/training_output/classify_run_gpu_fp16_0818_no_cnn/model.pt"

    MUTATION_DATA_PATH = "/gpfs/flash/home/yzq/project/model/SomaticMutaGPT/data/processed/combined_mutation_data_0909/combined_processed_mutations_base_features.parquet"
    DATA_ROOT_PATH = "/gpfs/flash/home/yzq/project/model/SomaticMutaGPT/data/processed/combined_mutation_data_0909"
    CLINICAL_DATA_PATH = "/gpfs/flash/home/yzq/project/model/SomaticMutaGPT/data/processed/combined_processed_clinical_data/combined_processed_clinical_samples.parquet"

    # UMAP 结果和嵌入的输出路径
    OUTPUT_PLOT_PATH_ALL_TYPES = "/gpfs/flash/home/yzq/project/model/SomaticMutaGPT/results/plots/0818_umap_mutation_embeddings_by_manual_mapped_cancer_type_all_samples_no_cnn.png"
    OUTPUT_PLOT_PATH_BLADDER_VS_OTHER = "/gpfs/flash/home/yzq/project/model/SomaticMutaGPT/results/plots/0818_umap_mutation_embeddings_bladder_vs_other_all_samples_no_cnn.png"
    OUTPUT_PLOT_PATH_MELANOMA_VS_OTHER = "/gpfs/flash/home/yzq/project/model/SomaticMutaGPT/results/plots/0818_umap_mutation_embeddings_melanoma_vs_other_all_samples_no_cnn.png"
    OUTPUT_EMBEDDINGS_PATH = "/gpfs/flash/home/yzq/project/model/SomaticMutaGPT/results/plots/0818_mutation_embeddings_with_sample_ids_all_samples_no_cnn.parquet"

    EVO2_EMBEDDINGS_DIR = "/gpfs/flash/home/yzq/project/model/MutaGPT/data/processed/final_merged_evo2_embeddings_20_parts"
    EVO2_EMBEDDING_LAYER = "blocks.29.mlp.l3"

    # --- 模型和数据参数 ---
    BATCH_SIZE = 64
    NUM_WORKERS = 8
    # mask_ratio 设为 0.0，因为是推理，不需要掩码
    # max_seq_len 应该与训练时使用的 max_seq_len 保持一致
    DATA_ARGS = {"mask_ratio": 0.0, "max_seq_len": 1024}

    model_config_path_for_task = Path(MODEL_PATH).parent / 'model_config.json'
    if model_config_path_for_task.exists():
        with open(model_config_path_for_task, "r") as f:
            model_config_for_task = json.load(f)
        TASK_FOR_EMBEDDING = model_config_for_task.get('task', 'pretrain') # 默认预训练
    else:
        TASK_FOR_EMBEDDING = 'classify' if 'classify' in str(MODEL_PATH) else 'pretrain'
        warnings.warn(f"未找到模型配置文件 {model_config_path_for_task}，将任务类型猜测为 '{TASK_FOR_EMBEDDING}'。")


     # 1. 实例化 EmbeddingExtractorModule
    embedding_model = EmbeddingExtractorModule(
        task=TASK_FOR_EMBEDDING,
        model_path=MODEL_PATH,
        data_root=DATA_ROOT_PATH
    )
    print("EmbeddingExtractorModule 实例化成功。")

    # 2. 使用 MutationDataModule 准备数据
    print(f"正在使用 MutationDataModule 准备数据...")
    data_module = MutationDataModule(
        df_path=MUTATION_DATA_PATH,
        data_root=DATA_ROOT_PATH,
        task=TASK_FOR_EMBEDDING,
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        pin_memory=True,
        shuffle=False,
        clinical_df_path=CLINICAL_DATA_PATH,
        col_label='clin_CANCER_TYPE_manual_mapped_id',
        col_sequencing_method='sequencing_method_str_id',
        col_exon_strand='exon_strand_id',
        bucket_batch_size=BATCH_SIZE,
        bucket_boundaries=None,
        stochastic_chrom_shuffling=False,
        random_unknown_strand_prob=0.0,
        evo2_embeddings_dir= EVO2_EMBEDDINGS_DIR, # 使用显式路径
        evo2_embedding_layer= EVO2_EMBEDDING_LAYER, # 使用显式路径
        data_args=DATA_ARGS
    )
    
    data_module.setup() # 仍然调用 setup 来加载 vocabularies 和处理 df_mutations
    df_all_mutations_for_embedding = pd.read_parquet(MUTATION_DATA_PATH)
    
    # 确保测序方法信息也合并到 df_all_mutations_for_embedding 中
    df_clinical_full_for_merge = pd.read_parquet(CLINICAL_DATA_PATH)[['sampleId', 'sequencing_method_str_id']].drop_duplicates(subset=['sampleId'])
    df_all_mutations_for_embedding = pd.merge(df_all_mutations_for_embedding, df_clinical_full_for_merge, on='sampleId', how='left')
    df_all_mutations_for_embedding['sequencing_method_str_id'].fillna(0, inplace=True) # 填充NaN
    df_all_mutations_for_embedding['sequencing_method_str_id'] = df_all_mutations_for_embedding['sequencing_method_str_id'].astype(int)

    all_samples_dataset = MutationPretrainDataset(
        df_all_mutations_for_embedding,
        data_root=DATA_ROOT_PATH,
        col_sequencing_method='sequencing_method_str_id',
        col_exon_strand='exon_strand_id',
        stochastic_chrom_shuffling=False,
        random_unknown_strand_prob=0.0,
        evo2_embeddings_dir=EVO2_EMBEDDINGS_DIR,
        evo2_embedding_layer=EVO2_EMBEDDING_LAYER,
        **DATA_ARGS # 传递 mask_ratio=0.0 和 max_seq_len
    )
    
    # 使用 LengthBasedBatchSampler 为所有样本创建 DataLoader
    all_samples_sampler = LengthBasedBatchSampler(
        all_samples_dataset,
        batch_size=BATCH_SIZE,
        bucket_boundaries=data_module.bucket_boundaries, # 使用 DataModule 的 bucket_boundaries
        drop_last=False,
        shuffle=False
    )
    all_samples_dataloader = DataLoader(
        all_samples_dataset,
        num_workers=NUM_WORKERS,
        pin_memory=True,
        worker_init_fn=worker_init_fn_seed,
        batch_sampler=all_samples_sampler,
        collate_fn=collate_fn_for_mutation_data
    )
    
    sample_ids_ordered = all_samples_dataset.sample_ids # 获取所有样本的ID
    print(f"数据集中总样本数 (用于嵌入生成): {len(sample_ids_ordered)}")

    # 3. 使用 PyTorch Lightning Trainer 运行推理
    print("正在为所有样本生成嵌入...")
    trainer = Trainer(
        accelerator='cuda' if torch.cuda.is_available() else 'cpu', # 明确使用 cuda
        devices=1, # 提取嵌入通常在单GPU上进行，如果需要多GPU，需要配置DDP
        logger=False,
        enable_checkpointing=False,
        enable_progress_bar=True,
    )

    trainer.test(embedding_model, dataloaders=all_samples_dataloader)
    embeddings = embedding_model.all_embeddings
    print(f"生成的嵌入形状: {embeddings.shape}")

    # 4. 创建嵌入 DataFrame 并保存
    print("正在创建嵌入 DataFrame...")
    df_embeddings = pd.DataFrame(embeddings, columns=[f'emb_{i}' for i in range(embeddings.shape[1])])
    # 确保 sample_ids_ordered 的长度与 embeddings 的行数匹配
    # 注意：如果使用 DDP 且 devices > 1，这里的 sample_ids_ordered 应该是当前 rank 的子集
    # 但由于我们设置 devices=1，所以这里是所有样本
    df_embeddings['sampleId'] = sample_ids_ordered[:embeddings.shape[0]]

    print(f"正在将嵌入保存到 {OUTPUT_EMBEDDINGS_PATH}...")
    # 确保输出目录存在
    Path(OUTPUT_EMBEDDINGS_PATH).parent.mkdir(parents=True, exist_ok=True)
    df_embeddings.to_parquet(OUTPUT_EMBEDDINGS_PATH, index=False)
    print("嵌入保存成功。")

    # 5. 加载临床数据
    print(f"正在从 {CLINICAL_DATA_PATH} 加载临床数据...")
    df_clinical = pd.read_parquet(CLINICAL_DATA_PATH)
    if 'sampleId' not in df_clinical.columns:
        if df_clinical.index.name == 'sampleId':
            df_clinical = df_clinical.reset_index()
        else:
            raise ValueError("临床 DataFrame 必须包含 'sampleId' 列或将 'sampleId' 作为索引。")
    print(f"临床数据加载完成。形状: {df_clinical.shape}")

    # 6. 与临床数据合并
    print("正在将嵌入与临床数据合并...")
    # 确保合并时只保留需要的临床列，避免内存浪费
    required_clinical_cols = ['sampleId', 'clin_CANCER_TYPE_manual_mapped', 'sequencing_method_str'] # 增加测序方法str

    # 过滤临床数据，只保留实际存在的列
    available_clinical_cols = [col for col in required_clinical_cols if col in df_clinical.columns]

    df_merged = pd.merge(df_embeddings, df_clinical[available_clinical_cols], on='sampleId', how='inner')

    print(f"合并后数据形状: {df_merged.shape}")

    initial_merged_rows = df_merged.shape[0]
    df_merged.dropna(subset=['clin_CANCER_TYPE_manual_mapped'], inplace=True)
    if df_merged.shape[0] < initial_merged_rows:
        print(f"因缺少 'clin_CANCER_TYPE_manual_mapped' 而移除了 {initial_merged_rows - df_merged.shape[0]} 行。")
    print(f"删除 NaN 归类癌症类型后合并数据形状: {df_merged.shape}")

    # --- 修正：使用新的标准化癌症类型名称 (从 data_preprocessing_clinical.py 确保一致) ---
    BLADDER_CANCER_STD_NAME = "Bladder & Urothelial Cancer"
    MELANOMA_STD_NAME = "Skin Cancer (incl. Melanoma)"

    df_merged['cancer_type_group_bladder'] = df_merged['clin_CANCER_TYPE_manual_mapped'].apply(
        lambda x: BLADDER_CANCER_STD_NAME if x == BLADDER_CANCER_STD_NAME else 'Other Cancer Types'
    )
    print(f"已创建 'cancer_type_group_bladder' 列用于 '{BLADDER_CANCER_STD_NAME}' 二分类绘图。")

    df_merged['cancer_type_group_melanoma'] = df_merged['clin_CANCER_TYPE_manual_mapped'].apply(
        lambda x: MELANOMA_STD_NAME if x == MELANOMA_STD_NAME else 'Other Cancer Types'
    )
    print(f"已创建 'cancer_type_group_melanoma' 列用于 '{MELANOMA_STD_NAME}' 二分类绘图。")
    # --- 修正结束 ---

    # 7. 执行 UMAP 降维
    print("正在执行 UMAP 降维...")
    reducer = umap.UMAP(
        random_state=42,
        n_components=2,
        n_neighbors=150, # 默认值，可根据数据量调整
        min_dist=0.1,    # 默认值，可调整
        metric='cosine'  # 对于高维嵌入通常是好的选择
    )

    embedding_cols = [col for col in df_merged.columns if col.startswith('emb_')]
    if not embedding_cols:
        raise ValueError("合并 DataFrame 中未找到嵌入列。请检查列命名。")

    embeddings_for_umap = df_merged[embedding_cols].values

    if embeddings_for_umap.shape[0] < 2:
        print("警告：样本数量不足以执行 UMAP。跳过 UMAP 和绘图。")
        return

    # 调整 n_neighbors 以适应较小的样本数量
    if embeddings_for_umap.shape[0] < reducer.n_neighbors:
        print(f"警告：样本数量 ({embeddings_for_umap.shape[0]}) 小于 UMAP 的默认 n_neighbors ({reducer.n_neighbors})。正在调整 n_neighbors。")
        reducer.n_neighbors = max(2, embeddings_for_umap.shape[0] - 1) # 至少为2

    umap_embeddings = reducer.fit_transform(embeddings_for_umap)
    df_merged['umap_x'] = umap_embeddings[:, 0]
    df_merged['umap_y'] = umap_embeddings[:, 1]
    print("UMAP 完成。")

    # 8. 生成并保存 UMAP 图
    print("正在生成 UMAP 图...")

    # --- 图 1: 所有手动映射的肿瘤类型注释 ---
    # 确保输出目录存在
    Path(OUTPUT_PLOT_PATH_ALL_TYPES).parent.mkdir(parents=True, exist_ok=True)
    plot_umap(
        df_data=df_merged,
        hue_column='clin_CANCER_TYPE_manual_mapped',
        title=f'UMAP of Mutation Embeddings by Top {len(df_merged["clin_CANCER_TYPE_manual_mapped"].unique())} Cancer Types (All Samples)',
        output_path=OUTPUT_PLOT_PATH_ALL_TYPES,
        palette=sns.color_palette("hls", len(df_merged['clin_CANCER_TYPE_manual_mapped'].unique()))
    )

    # --- 图 2: 膀胱癌 vs. 其他肿瘤类型注释 ---
    custom_palette_bladder = {BLADDER_CANCER_STD_NAME: 'red', 'Other Cancer Types': 'blue'}
    plot_umap(
        df_data=df_merged,
        hue_column='cancer_type_group_bladder',
        title=f'UMAP of Mutation Embeddings: {BLADDER_CANCER_STD_NAME} vs. Other Cancer Types (All Samples)',
        output_path=OUTPUT_PLOT_PATH_BLADDER_VS_OTHER,
        palette=custom_palette_bladder
    )

    # --- 图 3: 黑色素瘤 vs. 其他肿瘤类型注释 ---
    custom_palette_melanoma = {MELANOMA_STD_NAME: 'green', 'Other Cancer Types': 'purple'}
    plot_umap(
        df_data=df_merged,
        hue_column='cancer_type_group_melanoma',
        title=f'UMAP of Mutation Embeddings: {MELANOMA_STD_NAME} vs. Other Cancer Types (All Samples)',
        output_path=OUTPUT_PLOT_PATH_MELANOMA_VS_OTHER,
        palette=custom_palette_melanoma
    )

    # --- 新增图：按测序方法分组的UMAP图 ---
    if 'sequencing_method_str' in df_merged.columns:
        OUTPUT_PLOT_PATH_SEQ_METHOD = Path(OUTPUT_PLOT_PATH_ALL_TYPES).parent / "0818_umap_mutation_embeddings_by_sequencing_method_no_cnn.png"
        plot_umap(
            df_data=df_merged,
            hue_column='sequencing_method_str',
            title=f'UMAP of Mutation Embeddings by Sequencing Method (All Samples)',
            output_path=OUTPUT_PLOT_PATH_SEQ_METHOD,
            palette='tab10' # 使用不同的调色板
        )

    print("所有图已成功生成并保存。")

if __name__ == '__main__':
    generate_embeddings_and_umap_plot()